from typing import Dict, List
import copy

from .strategy_generator import StrategyGenerator,FollowingStrategyGenerator
from ..shard.placement_types import DTensorSpec,Placement,Replicate,PlacementStrategy,Shard,_Partial,ShardingStrategy
import itertools
__all__ = ['OutputGenerator']


class OutputGenerator(StrategyGenerator):
    """
    OutputGenerator is a generic class to generate strategies for Output Node.
    """
    def __init__(self, operation_data_mapping, device_mesh,pre_strategies_vectors):
        self.op_data = operation_data_mapping
        self.device_mesh = device_mesh
        self.pre_strategies_vectors = pre_strategies_vectors
    # 是否要约束最后的output是partial   TODO
    def collate_strategies(self) -> List[ShardingStrategy]:
        strategy_list = []
        for index, strategy in enumerate(self.pre_strategies_vectors):
            input_specs = []
            # print(strategy.sharding_specs.output_specs)
            pre_spec=strategy.sharding_specs.output_specs
            input_specs.append(DTensorSpec(mesh = self.device_mesh, placements=pre_spec.placements,tensor_meta=self.op_data['input_0'].data))
            output_specs = DTensorSpec(mesh = self.device_mesh, placements=(Replicate(),),tensor_meta=self.op_data['output'].data)
            sharding_specs = PlacementStrategy(input_specs=input_specs,output_specs=output_specs)
            strategy_list.append(ShardingStrategy(name=str(sharding_specs),sharding_specs=sharding_specs,compute_cost =0))
        #input_specs=[]
        #strategy_list = []
        #input_specs.append(DTensorSpec(mesh = self.device_mesh, placements=(Replicate(),),tensor_meta=self.op_data['input_0'].data))
        #output_specs = DTensorSpec(mesh = self.device_mesh, placements=(Replicate(),),tensor_meta=self.op_data['output'].data)
        #sharding_specs = PlacementStrategy(input_specs=input_specs,output_specs=output_specs)
        #name = str(sharding_specs)
        #strategy_list.append(ShardingStrategy(name=name,sharding_specs=sharding_specs,compute_cost =0))

        return strategy_list
